import copy
import random

import numpy as np
import elkai
import threading
import matplotlib.pyplot as plt

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)


def create_data_model(file_path):
    raw = np.loadtxt(file_path, skiprows=6, encoding='utf-8', dtype=float)
    res_data = dict()
    res_data['num_indexes'] = len(raw)
    res_data['x'] = raw[:, 1]
    res_data['y'] = raw[:, 2]
    i = np.arange(len(raw))
    j = i.reshape(len(i), -1)
    res_data['distance_matrix'] = np.sqrt((raw[:, 1][i] - raw[:, 1][j]) ** 2 + (raw[:, 2][i] - raw[:, 2][j]) ** 2)
    res_data['num_vehicles'] = 1
    res_data['depot'] = 0
    return res_data


class Solver(threading.Thread):
    def __init__(self, threadID, name, indexes, matrix):
        super().__init__()
        self.threadId = threadID
        self.name = name
        self.indexes = indexes
        self.matrix = matrix
        self.result = None

    def run(self):
        result = elkai.solve_float_matrix(matrix=self.matrix)
        self.result = self.indexes[result]


def evaluate(assignment, matrix):
    solvers = []
    for _m in range(m):
        indexes = np.argwhere(assignment == _m + 1).squeeze()
        indexes = np.insert(indexes, 0, 0)
        sub_matrix = matrix[indexes, :][:, indexes]
        solver = Solver(threadID=_m, name='route-1', indexes=indexes, matrix=sub_matrix)
        solver.start()
        solvers.append(solver)

    for thread in solvers:
        thread.join()

    total = 0
    routes = []
    for i, solver in enumerate(solvers):
        route = solver.result.tolist()
        routes.append(route)
        length = sum(matrix[route[j - 1]][route[j]] for j in range(1, len(route))) + matrix[route[-1]][route[0]]
        total += length

    return total, routes


def roulette_selection(weights, num_winners):
    weights = np.array(weights)
    weights = 1.0 / weights
    total = sum(weights)
    weights /= total
    return np.random.choice(range(len(weights)), size=num_winners, replace=True, p=weights)


def crossover(p1, p2):
    point = np.random.randint(len(p1))
    p1[point:], p2[point:] = p2[point:], p1[point:]
    return p1, p2


m = 3
NP = 50
Pc = 0.7
Pm = 0.2
MAX_GEN = 1000
plt.figure(figsize=(10, 4))
if __name__ == '__main__':
    data = create_data_model('../../data/TSPLIB/eil51.tsp.txt')
    # Initialization
    assignments = np.random.randint(low=1, high=m + 1, size=(NP, data['num_indexes']))  # random assignments
    assignments[:, 0] = 0  # set 0 to all salesmen
    # Evolution Loop
    best = 1e99
    best_assignment = None
    history = []
    for generation in range(MAX_GEN):
        # Evaluation
        fitness = [evaluate(x, matrix=data['distance_matrix'])[0] for x in assignments]
        if best > min(fitness):
            best = min(fitness)
            best_assignment = copy.deepcopy(assignments[fitness.index(min(fitness))])
        history.append(best)

        plt.clf()
        plt.subplot(1, 2, 1)
        plt.plot(range(len(history)), history)
        plt.title(f'{generation}: {history[-1]:.2f}')
        plt.subplot(1, 2, 2)
        length, routes = evaluate(best_assignment, matrix=data['distance_matrix'])
        for route in routes:
            plt.plot(data['x'][route + [route[0]]], data['y'][route + [route[0]]])
        plt.title(f'{generation}: {length:.2f}')
        plt.pause(0.001)

        # Selection
        min_idx = fitness.index(min(fitness))  # reserve the best
        optimal_assignment = copy.deepcopy(assignments[min_idx])
        winners = roulette_selection(weights=fitness, num_winners=NP)
        winners[min_idx] = min_idx
        assignments = assignments[winners]
        # Crossover
        count = 0
        for i in range(NP):
            if random.uniform(0, 1) < Pc:
                j = i
                count += 1
                if count == 2:
                    assignments[i], assignments[j] = crossover(assignments[i], assignments[j])
                    count = 0
        # Mutation
        for i in range(NP):
            prob = np.random.uniform(size=len(assignments[i]))
            rand = np.random.randint(low=1, high=m + 1, size=len(assignments[i]))
            assignments[i][prob < Pc] = rand[prob < Pc]
            assignments[i][0] = 0

    # route = elkai.solve_float_matrix(matrix=data['distance_matrix'])
    # route_length = sum(data['distance_matrix'][route[i - 1]][route[i]] for i in range(1, len(route))) + \
    #                data['distance_matrix'][route[-1]][route[0]]
    # plt.plot(data['x'][route + [route[0]]], data['y'][route + [route[0]]])
    # plt.scatter(data['x'], data['y'])
    # for idx, (x, y) in enumerate(zip(data['x'], data['y'])):
    #     plt.text(x, y, s=idx)
    # plt.title(f'{route_length:.2f}')
    # plt.show()
